
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def create_custom_vector_field():
    """
    创建自定义向量场数据
    """
    # 定义更精细的网格
    x = np.linspace(-2, 2, 25)
    y = np.linspace(-2, 2, 25)
    z = np.linspace(-2, 2, 25)
    
    X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
    
    # 创建更复杂的向量场
    U = -Y * np.exp(-0.1*(X**2 + Y**2 + Z**2))
    V = X * np.exp(-0.1*(X**2 + Y**2 + Z**2))
    W = 0.5 * Z * np.exp(-0.1*(X**2 + Y**2 + Z**2))
    
    return X, Y, Z, U, V, W

def plot_advanced_streamlines(x, y, z, u, v, w):
    """
    高级流线图绘制函数
    """
    fig = plt.figure(figsize=(15, 12))
    ax = fig.add_subplot(111, projection='3d')
    
    # 计算速度大小用于着色
    speed = np.sqrt(u**2 + v**2 + w**2)
    
    # 绘制流线图
    stream = ax.streamplot(x, y, z, u, v, w,
                              linewidth=2,
                              cmap='plasma',
                              arrowstyle='-|>',
                              arrowsize=1.5,
                              density=1.5,
                              color=speed)
    
    # 设置坐标轴
    ax.set_xlabel('X坐标', fontsize=14)
    ax.set_ylabel('Y坐标', fontsize=14)
    ax.set_zlabel('Z坐标', fontsize=14)
    ax.set_title('高级三维流线图可视化', fontsize=16)
    
    # 添加颜色条
    cbar = plt.colorbar(stream.lines, ax=ax)
    cbar.set_label('流速强度', fontsize=12)
    
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    # 使用自定义数据
    x, y, z, u, v, w = create_custom_vector_field()
    plot_advanced_streamlines(x, y, z, u, v, w)
